package com.xl.bigdata.es.index;

import com.xl.bigdata.es.xpack.ESXpackRestHighLevelClient;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.UpdateByQueryRequest;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;

import java.io.IOException;
import java.util.Collections;
import java.util.List;

/**
 * @ClassName UpdateByQueryAPI
 * @Description TODO
 * @Author xl.gao
 * @Date 2022/1/4 10:32
 * @Version
 */
public class UpdateByQueryAPI {

    public static boolean updateByQueryAPI(RestHighLevelClient client, String index) {

        UpdateByQueryRequest updateByQueryRequest = new UpdateByQueryRequest(index);
        updateByQueryRequest.setQuery(new TermQueryBuilder("citycomple", "广州市"));
        updateByQueryRequest.setBatchSize(1000);

        Script painless = new Script(ScriptType.INLINE, "painless",
                "if(ctx._source.districtcomple == 'thq') {ctx._source.districtcomple = '天河区'}",
                Collections.emptyMap());

        updateByQueryRequest.setScript(painless);
        updateByQueryRequest.setSlices(12);  // 1：8178 ， 2:6000
        updateByQueryRequest.setScroll(TimeValue.timeValueMinutes(10));
        updateByQueryRequest.setTimeout(TimeValue.timeValueMinutes(2));
        updateByQueryRequest.setRefresh(true);

        try {
            BulkByScrollResponse bulkByScrollResponse = client.updateByQuery(updateByQueryRequest, RequestOptions.DEFAULT);

            TimeValue took = bulkByScrollResponse.getTook();
            long total = bulkByScrollResponse.getTotal();
            long updated = bulkByScrollResponse.getUpdated();
            System.out.println("响应时间：(ms)" + took.getMillis() + " 总量：" + total + " 更新数量：" + updated);
            List<BulkItemResponse.Failure> bulkFailures = bulkByScrollResponse.getBulkFailures();
            if (bulkFailures != null && bulkFailures.size() > 0) {
                System.out.println("失败任务数量：" + bulkFailures.size());
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        return false;
    }

    public static void main(String[] args) {

        RestHighLevelClient restHighLevelClient = ESXpackRestHighLevelClient.getRestHighLevelClient();


        updateByQueryAPI(restHighLevelClient, "address_base_02");

        ESXpackRestHighLevelClient.closeRestHighLevelClient();
    }
}
